"""Module for the Lagrangian flux functions."""

from hogs.solvers.flux.flux_functions import FluxFunction
from hogs.solvers.riemann.euler import reuler
from hogs.sph.sph import SPHInterpolate, NNPS

import numpy

###########################################################################
#`LxfFluxLagrangian` class
###########################################################################
class LxfFluxLagrangian(FluxFunction):
    """Lagrangian version of the Lax Fredrichs flux described in

    'A high order ENO conservative Lagrangian type scheme for the compressible
    Euler equations.' by Juan Cheng and Chi-Wang Shu, JCP 227, 1567-1596

    """
    
    def compute_flux(self, grid, k=None):

        # iterate over the cell faces at (i-1/2)
        for face in grid.faces[grid.nb:-grid.nb]:

            cl, cr = face.get_sound_speeds()
            alpha = max(cl, cr)
            
            # get the conserved and primitive variables adjoinng this face
            wl, wr = face.get_primitive_variables()
            ql, qr = face.get_conserved_variables()

            # get the Roe's average velocity at the face
            xidot = face.compute_roe_average( wl[1], wr[1] )

            # The Roe's average velocity is the vertex velocity
            face.vertex_vel = xidot

            # compute the flux
            f = face.flux

            f[0] = -0.5 * alpha * (qr[0] - ql[0])

            f[1] = 0.5 * ( (wl[2] + wr[2]) - alpha * (qr[1] - ql[1]) )

            f[2] = 0.5 * ( (wl[1]*wl[2] + wr[1]*wr[2]) - \
                           alpha * (qr[2] - ql[2]) )
            
###########################################################################
#`GodunovFluxLagrangian` class
###########################################################################
class GodunovFluxLagrangian(FluxFunction):
    """Lagrangian version of the Godunov's flux defined in.

    'A high order ENO conservative Lagrangian type scheme for the compressible
    Euler equations.' by Juan Cheng and Chi-Wang Shu, JCP 227, 1567-1596

    """
    def __init__(self, gamma=1.4):
        self.gamma = gamma
        self.rsolver = reuler.RiemannSolverEulerExact(gamma)

    def compute_flux(self, grid, dt):

        nb = grid.nb

        # compute the fluxes
        for face in grid.faces[nb:-nb]:

            # get the primitive variables adjoining the face
            wl, wr = face.get_primitive_variables()

            # set the flux
            self.rsolver.set_lagrangian_godunov_flux(wl, wr, face)

##########################################################################
#`SPHFluxLagrangian` class
########################################################################## 
class SPHFluxLagrangian(FluxFunction):
    """Flux evaluation using the SPH interpolation technique.

    The Lagrangian pressure flux F (0, p, pu) is estimated at the cell
    interfaces (which define the instantaneous volumes). These
    interpolated quantities define the flux at the interfaces.

    """
    def __init__(self):
        self.nnps = NNPS()    

    def compute_flux(self, grid, k=None):

        nnps = self.nnps
        x = grid.x

        # set a reasonable smoothing length parameter from cell sizes
        #dx = max( x[1:] - x[:-1] )
        #h = 2*dx

        # bin the points 
        #nnps.bin(grid, cell_size = 2 * h)

        # get the density of data points
        #rho = SPHInterpolate.get_density(nnps, grid.xc, h)

        # get the SPH interpolated pressure and velocity
        #u = SPHInterpolate.interpolate(nnps, grid.x, grid.xc, grid.w[1], rho, h)
        #p = SPHInterpolate.interpolate(nnps, grid.x, grid.xc, grid.w[2], rho, h)

        for face in grid.faces[grid.nb:-grid.nb]:
        
            cl, cr = face.get_sound_speeds()
            alpha1 = max(cl, cr)

            # get the primitive variables adjoinng this face
            wl, wr = face.get_primitive_variables()
            
            # get the Roe's average velocity at the face
            xidot = face.compute_roe_average( wl[1], wr[1] )
            
            # The Roe averaged velocity is the vertex velocity
            face.vertex_vel = xidot

            # get the SPH interpolated pressure and velocity
            #_p = p[ face.index ]
            #_u = u[ face.index ]

            # compute some averaged state
            wlm1, wrm1 = face.get_primitive_variables(i=-1)
            alpha2 = max( face.get_sound_speeds(i=-1) )
            
            wlp1, wrp1 = face.get_primitive_variables(i=+1)
            alpha3 = max( face.get_sound_speeds(i=+1) )

            _u = 0.1 * wlm1[1] + 0.4 * (wl[1] + wr[1]) + 0.1 * wrp1[1]
            _p = 0.1 * wlm1[2] + 0.4 * (wl[2] + wr[2]) + 0.1 * wrp1[2]

            # store the fluxes
            f = face.flux

            ql, qr = face.get_conserved_variables()
            dq = 0.5 * max( alpha1, alpha2, alpha3 ) * (qr - ql)

            f[0] = 0.0 - dq[0]
            f[1] = _p  - dq[1]
            f[2] = _p * _u - dq[2]

flux_types = [

    ("LaxFredrichs", LxfFluxLagrangian),

    ("Godunov", GodunovFluxLagrangian),

    ("SPH", SPHFluxLagrangian)

    ]
